# 弹簧阻尼系统的LQR控制:跟踪问题：解决存在稳态误差的问题
# 输入和稳态误差权衡的问题：输入小，无法达到平衡状态，稳态误差大，系统会变得不稳定
# xd[k+1] = ADXd[k],AD目标转移矩阵，为单位阵时，追踪的目标为常数

from datetime import datetime
import F_LQR_Gain 
import numpy as np
import math
import matplotlib.pyplot as plt
import pandas as pd
import control as ct
import time
import scipy.linalg as la
from scipy.signal import StateSpace
import F2_InputAugmentMatrix_SS_U

# 1.定义系统
m = 1
k = 1
b = 0.5
# x[k+1] = Ax[k] + Bu[k]
A = np.array([[0,1],[-k/m,-b/m]])
B = np.array([[0],[1/m]])
C = np.array([[1,0]])
D = 0
# 创建状态空间模型
# sys = StateSpace(A, B, C, D)
Ts = 0.1
sys_continuous = ct.StateSpace(A, B, C, D)
sys_discrete = ct.sample_system(sys_continuous, Ts, method='zoh')
# 获取行数
n = len(A)
# 获取列数
p =  B.shape[1]
# 离散化后的状态矩阵和输入矩阵
A = sys_discrete.A
B = sys_discrete.B
# 目标
xd = np.array([[1],[0]]).reshape(1,-1)
AD = np.eye(n)
# 2.初始化系统
x0 = np.array([[0],[0]]).reshape(1,-1)
x = x0
xa = np.hstack((x,xd))
# 输入
u0 = np.array([[0]]).reshape(1,-1)
u = u0

# 定义系统运行步数
k_steps = 100
x_history = np.zeros((n,k_steps+1))
x_history[:,0] = x
u_history = np.zeros((p,k_steps))
u_history[:,0] = u

# 设置权重
Q = np.array([[1,0],[0,1]])
S = np.array([[1,0],[0,1]])
R=10
N = k_steps
P_k = S
# 组合新的矩阵
Aa,Ba,Qa,Sa,R,ud = F2_InputAugmentMatrix_SS_U.F2_InputAugmentMatrix_SS_U(A,B,Q,R,S,xd)
# 3.计算反馈增益
F = F_LQR_Gain.LQR_Gain(Aa,Ba,Qa,R,Sa)

for k in range(1,k_steps+1):
    u =-F@xa.T + ud
    x=A@x_history[:,k-1].reshape(-1,1)+B@u
    xa = np.vstack((x,xd.T)).T
    x_history[:,k]=x.T
    u_history[:,k-1]=u
"""
# # 写入数据到表格中
from openpyxl import Workbook
# 创建 Excel 工作簿
wb = Workbook()
ws = wb.active
# 将矩阵数据写入工作表
for row in x_history:
    ws.append(row.tolist())
# 保存 Excel 文件
wb.save('x_history.xlsx')
wb = Workbook()
ws = wb.active

# 将矩阵数据写入工作表
for row in u_history:
    ws.append(row.tolist())
# 保存 Excel 文件
wb.save('u_history.xlsx')

# 创建 Excel 工作簿
wb = Workbook()
ws = wb.active
# 将矩阵数据写入工作表
for row in F_N:
    ws.append(row.tolist())
# 保存 Excel 文件
wb.save('F_N.xlsx')
from openpyxl import Workbook
wb = Workbook()
ws = wb.active
for row in u_history:
    ws.append(row.tolist())
# 保存 Excel 文件
wb.save('u_history.xlsx')
""" 

# 4.制图
plt.figure()
plt.subplot(2,1,1)
plt.title('Input and State Plot')
plt.xlabel('steps')
plt.ylabel('State')
for i in range(0,n): 
    plt.plot(x_history[i,:],label=f'State {i+1}')
    
plt.legend()
plt.subplot(2,1,2)
plt.ylabel('Input')
for i in range(0,p):
    plt.plot(u_history[i,:],label=f'Input {i+1}')
    
plt.legend()
plt.grid()
plt.show()